package cn.edu.hitsz.compiler.asm;

import cn.edu.hitsz.compiler.NotImplementedException;
import cn.edu.hitsz.compiler.ir.IRImmediate;
import cn.edu.hitsz.compiler.ir.IRValue;
import cn.edu.hitsz.compiler.ir.IRVariable;
import cn.edu.hitsz.compiler.ir.Instruction;
import cn.edu.hitsz.compiler.utils.FileUtils;

import java.util.*;


/**
 * TODO: 实验四: 实现汇编生成
 * <br>
 * 在编译器的整体框架中, 代码生成可以称作后端, 而前面的所有工作都可称为前端.
 * <br>
 * 在前端完成的所有工作中, 都是与目标平台无关的, 而后端的工作为将前端生成的目标平台无关信息
 * 根据目标平台生成汇编代码. 前后端的分离有利于实现编译器面向不同平台生成汇编代码. 由于前后
 * 端分离的原因, 有可能前端生成的中间代码并不符合目标平台的汇编代码特点. 具体到本项目你可以
 * 尝试加入一个方法将中间代码调整为更接近 risc-v 汇编的形式, 这样会有利于汇编代码的生成.
 * <br>
 * 为保证实现上的自由, 框架中并未对后端提供基建, 在具体实现时可自行设计相关数据结构.
 *
 * @see AssemblyGenerator#run() 代码生成与寄存器分配
 */
public class AssemblyGenerator {
    /**
     * 中间代码指令
     */
    private List<Instruction> list;
    /**
     * 变量最后一次出现的位置
     */
    private Map<IRVariable,Integer> lastAppearance;
    /**
     * 寄存器的占用信息,未被占用的寄存器对应null
     */
    private Map<Reg,IRVariable> regs;
    /**
     * 生成的汇编指令
     */
    private List<Assembly> assemblyList;
    /**
     * 加载前端提供的中间代码
     * <br>
     * 视具体实现而定, 在加载中或加载后会生成一些在代码生成中会用到的信息. 如变量的引用
     * 信息. 这些信息可以通过简单的映射维护, 或者自行增加记录信息的数据结构.
     *
     * @param originInstructions 前端提供的中间代码
     */
    public void loadIR(List<Instruction> originInstructions) {
        // TODO: 读入前端提供的中间代码并生成所需要的信息
        list = originInstructions;
        //使用哈希表保存后续信息引用情况
        lastAppearance = new HashMap<>();
        getLastAppearance(list);
    }

    /**
     * 记录语法变量最后一次出现的位置
     * @param list 中间代码指令列表
     */
    private void getLastAppearance(List<Instruction> list){
        for (int i = 0; i < list.size(); i ++){
            Instruction instruction = list.get(i);
            switch(instruction.getKind()){
                case MOV -> {
                    IRVariable irVariable = instruction.getResult();
                    lastAppearance.put(irVariable,i);
                    IRValue irValue = instruction.getFrom();
                    if(irValue.isIRVariable()){
                        lastAppearance.put((IRVariable) irValue,i);
                    }
                }
                case RET -> {
                    IRValue irValue = instruction.getReturnValue();
                    if(irValue.isIRVariable()){
                        lastAppearance.put((IRVariable) irValue,i);
                    }
                }
                case ADD,SUB,MUL -> {
                    IRVariable result = instruction.getResult();
                    lastAppearance.put(result,i);
                    IRValue lhs = instruction.getLHS(), rhs = instruction.getRHS();
                    if(lhs.isIRVariable()){
                        lastAppearance.put((IRVariable) lhs,i);
                    }
                    if(rhs.isIRVariable()){
                        lastAppearance.put((IRVariable) rhs,i);
                    }
                }
            }
        }
    }

    /**
     * 初始化register哈希表，即寄存器初始占用信息
     */
    private void initialRegister(){
        Reg[] values = Reg.values();
        regs = new LinkedHashMap<>();
        for (Reg value : values){
            regs.put(value,null);
        }
        //排除a0寄存器，使之专用于返回结果
        regs.remove(Reg.a0);
    }

    /**
     * 为翻译过程的变量分配寄存器
     * @param irVariable 新的语法变量
     * @param index 当前语句索引序号
     */
    private Reg setReg(IRVariable irVariable, int index){
        for (Reg reg : regs.keySet()){
            IRVariable temp = regs.get(reg);
            if(temp == null || lastAppearance.get(temp) < index){
                regs.put(reg,irVariable);
                return reg;
            }
        }
        throw new RuntimeException("Registers are not enough!");
    }

    /**
     * 获取已经在寄存器里的变量的位置
     * @param irVariable 语法变量
     */
    private Reg getReg(IRVariable irVariable){
        for (Reg reg : regs.keySet()){
            IRVariable temp = regs.get(reg);
            if(temp != null && temp.equals(irVariable)) return reg;
        }
        return null;
    }

    /**
     * 执行代码生成.
     * <br>
     * 根据理论课的做法, 在代码生成时同时完成寄存器分配的工作. 若你觉得这样的做法不好,
     * 也可以将寄存器分配和代码生成分开进行.
     * <br>
     * 提示: 寄存器分配中需要的信息较多, 关于全局的与代码生成过程无关的信息建议在代码生
     * 成前完成建立, 与代码生成的过程相关的信息可自行设计数据结构进行记录并动态维护.
     */
    public void run() {
        // TODO: 执行寄存器分配与代码生成
        //throw new NotImplementedException();
        initialRegister();
        assemblyList = new ArrayList<>();
        for (int i = 0; i < list.size(); i ++){
            Instruction instruction = list.get(i);
            switch (instruction.getKind()){
                // IR 是 RET 指令
                case RET -> assemblyList.add(Assembly.createAddi(Reg.a0,
                        getReg((IRVariable)instruction.getReturnValue()),0));
                //IR 是 一元的，(有返回值, 有一个参数)
                case MOV -> {
                    //给返回值分配寄存器
                    Reg rd = getReg(instruction.getResult());
                    if(rd == null){
                        rd = setReg(instruction.getResult(), i);
                    }
                    IRValue irValue = instruction.getFrom();
                    if (irValue.isImmediate()) {
                        assemblyList.add(Assembly.createLi(rd, ((IRImmediate) irValue).getValue()));
                    } else if (irValue.isIRVariable()) {
                        assemblyList.add(Assembly.createAddi(rd, getReg((IRVariable) irValue),0));
                    }
                }
                case ADD,SUB,MUL -> {
                    Reg rd = getReg(instruction.getResult());
                    if(rd == null){
                        rd = setReg(instruction.getResult(), i);
                    }
                    IRValue lhs = instruction.getLHS();
                    IRValue rhs = instruction.getRHS();
                    switch (instruction.getKind()){
                        case ADD -> {
                            if (lhs.isImmediate() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) lhs).getValue()+((IRImmediate) rhs).getValue()));
                            } else if (lhs.isImmediate() && rhs.isIRVariable()) {
                                assemblyList.add(Assembly.createAddi(rd, getReg((IRVariable) rhs), ((IRImmediate) lhs).getValue()));
                            } else if (lhs.isIRVariable() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createAddi(rd, getReg((IRVariable) lhs), ((IRImmediate) rhs).getValue()));
                            } else {
                                assemblyList.add(Assembly.createAdd(rd, getReg((IRVariable) lhs), getReg((IRVariable) rhs)));
                            }
                        }
                        case SUB -> {
                            if (lhs.isImmediate() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) lhs).getValue() - ((IRImmediate) rhs).getValue()));
                            } else if (lhs.isImmediate() && rhs.isIRVariable()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) lhs).getValue()));
                                assemblyList.add(Assembly.createSub(rd, rd, getReg((IRVariable) rhs)));
                            } else if (lhs.isIRVariable() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createAddi(rd, getReg((IRVariable) lhs), -((IRImmediate) rhs).getValue()));
                            } else {
                                assemblyList.add(Assembly.createSub(rd, getReg((IRVariable) lhs), getReg((IRVariable) rhs)));
                            }
                        }
                        case MUL -> {
                            if (lhs.isImmediate() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) lhs).getValue() * ((IRImmediate) rhs).getValue()));
                            } else if (lhs.isImmediate() && rhs.isIRVariable()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) lhs).getValue()));
                                assemblyList.add(Assembly.createMul(rd, rd, getReg((IRVariable) rhs)));
                            } else if (lhs.isIRVariable() && rhs.isImmediate()) {
                                assemblyList.add(Assembly.createLi(rd, ((IRImmediate) rhs).getValue()));
                                assemblyList.add(Assembly.createMul(rd, rd, getReg((IRVariable) lhs)));
                            } else {
                                assemblyList.add(Assembly.createMul(rd, getReg((IRVariable) lhs), getReg((IRVariable) rhs)));
                            }
                        }
                    }
                }
            }
        }
    }

    public List<Assembly> getAssembly() {
        return this.assemblyList;
    }

    /**
     * 输出汇编代码到文件
     *
     * @param path 输出文件路径
     */
    public void dump(String path) {
        // TODO: 输出汇编代码到文件
        FileUtils.writeLines(path, getAssembly().stream().map(Assembly::toString).toList());
    }
}

